# -*- coding: utf-8 -*-
# ---------------------------------------------------------------------------
# 将TrainSet中所有单词按词频标成id，将id映射成一个长度为D的向量，(B, T, D)作为输入
# LSTM隐藏层同样为D，经过若干个LSTM后，输出为(B, T, D)，进行softmax后为(B, T, Vocab_size)
# 真实值为输入Tensor后移一个Timestep，损失用cross entropy计算，然后在B上平均，在T上求和
# ---------------------------------------------------------------------------
import tensorflow as tf
import numpy as np
import time
import loaddata
from ptb import *

# ---------------------------------------------------------------------------
# Read settings from cmd
# e.g.
# python main.py --data_path=data --save_path=tmp/bn --mode=bn
# ---------------------------------------------------------------------------
flags = tf.flags
flags.DEFINE_string("data_path", None,
                    "Where the training/test data is stored.")
flags.DEFINE_string("save_path", None,
                    "Model output directory.")
flags.DEFINE_string("mode", "std",
                    "LSTM mode of std, bn and bren.")
FLAGS = flags.FLAGS

def run_epoch(session, model, eval_op=None, verbose=False):
    """Runs the model on the given data."""
    start_time = time.time()
    costs = 0.0
    iters = 0
    state = session.run(model.initial_state)

    fetches = {
            "cost": model.cost,
            "final_state": model.final_state,
    }
    if eval_op is not None:
        fetches["eval_op"] = eval_op

    for step in range(model.input.epoch_size):
        feed_dict = {}

        # use last epoch's final state as new epoch's initial
        for i, (t, c, h) in enumerate(model.initial_state):
            # feed_dict[t] = state[i].t
            feed_dict[c] = state[i].c
            feed_dict[h] = state[i].h

        vals = session.run(fetches, feed_dict)
        cost = vals["cost"]
        state = vals["final_state"]

        costs += cost
        iters += model.input.num_steps

        if verbose and step % (model.input.epoch_size // 10) == 10:
            print("%.3f perplexity: %.3f speed: %.0f wps" %
                        (step * 1.0 / model.input.epoch_size, np.exp(costs / iters),
                         iters * model.input.batch_size / (time.time() - start_time)))

    return np.exp(costs / iters)

def main(_):
    # Prepare data
    # return list of ids and # of training vacab
    if not FLAGS.data_path:
        raise ValueError("Must set --data_path to PTB data directory")
    raw_data = loaddata.raw_data(FLAGS.data_path)
    train_data, valid_data, test_data, _ = raw_data

    # Prepare params
    config = get_config()
    eval_config = get_config()
    eval_config.batch_size = len(test_data) // (eval_config.num_steps + 1)
    # eval_config.num_steps = 1

    with tf.Graph().as_default():
        # ----------------------------------------------------------------------------------
        # Initial Model
        # ----------------------------------------------------------------------------------
        initializer = tf.random_uniform_initializer(-config.init_scale,
                                                     config.init_scale)
        with tf.name_scope("Train"):
            # Prepare feeding data and info
            train_input = PTBInput(config=config, data=train_data, name="TrainInput")
            # Build Model
            with tf.variable_scope("Model", reuse=None, initializer=initializer):
                m = PTBModel(is_training=True, config=config, input_=train_input,
                    mode=FLAGS.mode)
            # Prepare summary
            tf.summary.scalar("Training Loss", m.cost)
            tf.summary.scalar("Training Acc", m.acc)
            tf.summary.scalar("Learning Rate", m.lr)
        with tf.name_scope("Valid"):
            # Prepare feeding data and info
            valid_input = PTBInput(config=config, data=valid_data, name="ValidInput")
            # Build Model
            with tf.variable_scope("Model", reuse=True, initializer=initializer):
                mvalid = PTBModel(is_training=False, config=config, input_=valid_input,
                    mode=FLAGS.mode)
            # Prepare summary
            tf.summary.scalar("Validation Loss", mvalid.cost)
            tf.summary.scalar("Validation Acc", mvalid.acc)
        with tf.name_scope("Test"):
            # Prepare feeding data and info
            test_input = PTBInput(config=eval_config, data=test_data, name="TestInput")
            # Build Model
            with tf.variable_scope("Model", reuse=True, initializer=initializer):
                mtest = PTBModel(is_training=False, config=eval_config, input_=test_input,
                    mode=FLAGS.mode)

        # ----------------------------------------------------------------------------------
        # Start to train
        # ----------------------------------------------------------------------------------
        sv = tf.train.Supervisor(logdir=FLAGS.save_path)
        with sv.managed_session() as session:
            for i in range(config.max_max_epoch):
                lr_decay = config.lr_decay ** max(i + 1 - config.max_epoch, 0.0)
                m.assign_lr(session, config.learning_rate * lr_decay)
                print("\nEpoch: %d Learning rate: %.3f Training: %s"
                    % (i + 1, session.run(m.lr), m.training))
                train_perplexity = run_epoch(session, m, eval_op=m.train_op,
                                                                         verbose=True)
                print("Epoch: %d Train Perplexity: %.3f" % (i + 1, train_perplexity))
                
                print("\nTraining: %s" % mvalid.training)
                valid_perplexity = run_epoch(session, mvalid)
                print("Epoch: %d Valid Perplexity: %.3f" % (i + 1, valid_perplexity))

            print("\nTraining: %s" % mtest.training)
            test_perplexity = run_epoch(session, mtest)
            print("Test Perplexity: %.3f" % test_perplexity)

            if FLAGS.save_path:
                print("Saving model to %s." % FLAGS.save_path)
                sv.saver.save(session, FLAGS.save_path, global_step=sv.global_step)

if __name__ == "__main__":
    tf.app.run()